#!/usr/bin/env python3
import sys
from scapy.all import Ether, sendp, get_if_hwaddr, get_if_list, IP, ICMP
from scapy.packet import Packet, bind_layers
from scapy.fields import ByteField

# 自定义 SR EtherType
TYPE_SR = 0x0908

# ---------------------- 定义 SourceRouting 头部 ----------------------
class SourceRouting(Packet):
    name = "SourceRouting"
    fields_desc = [
        ByteField("path_len", 0),      # 1字节 path_len
        ByteField("current_index", 0), # 1字节 current_index
        ByteField("port0", 0),         # 必须4个端口字段
        ByteField("port1", 0),
        ByteField("port2", 0),
        ByteField("port3", 0)
    ]

# 绑定 Ether 类型到自定义头部
bind_layers(Ether, SourceRouting, type=TYPE_SR)

# ---------------------- 获取网络接口 ----------------------
def get_iface():
    for iface in get_if_list():
        if "eth0" in iface:
            return iface
    raise RuntimeError("Cannot find eth0 interface")

# ---------------------- 主函数 ----------------------
def main():
    if len(sys.argv) < 3:
        print("Usage: send_sr.py <dst_ip> <port0> [<port1> <port2> <port3>]")
        print("Example: send_sr.py 10.0.2.2 4 2 1 0")  # 注意必须提供4个端口参数
        sys.exit(1)

    dst_ip = sys.argv[1]
    # 读取端口参数并补齐4个（不足补0）
    input_ports = list(map(int, sys.argv[2:]))
    path_ports = input_ports + [0] * (4 - len(input_ports))

    iface = get_iface()
    src_mac = get_if_hwaddr(iface)
    dst_mac = "08:00:00:00:01:00"  # 目标 MAC（交换机入口端口）

    # 构造源路由头部（关键修复）
    sr_hdr = SourceRouting(
        path_len=len([p for p in path_ports if p != 0]),  # 实际路径长度
        current_index=0,
        port0=path_ports[0],
        port1=path_ports[1],
        port2=path_ports[2],
        port3=path_ports[3]
    )

    # 构造完整数据包（包含 IP/ICMP 用于验证）
    pkt = Ether(src=src_mac, dst=dst_mac, type=TYPE_SR) / sr_hdr / IP(src="10.0.1.1", dst=dst_ip) / ICMP()

    print(f"Sending SR packet on {iface}: path_ports={path_ports}")
    pkt.show2()  # 验证头部结构
    sendp(pkt, iface=iface, verbose=False)

if __name__ == "__main__":
    main()
